#ifndef __SMIMEMESSAGE_HPP__
#define __SMIMEMESSAGE_HPP__

#include <string>
#include <openssl/pkcs7.h>

#include "GenericException.hpp"
#include "SSLWrapper.hpp"
#include "Logger.hpp"

/**
 * class represents an S/MIME message and provides functions for decryption and signature checking
 */
class SMIMEMessage {
	private:
		PKCS7 *PKCS;
		BIO *Content;
		BIO *Decrypted;
		BIO *Verified;

		/**
		 * internal function
		 */
		std::string ReadFromBIO(BIO *f) {
			std::string Result;
			char Buffer[256];
			while (BIO_gets(f, Buffer, sizeof(Buffer)) > 0) {
				Result += Buffer;
			}
			return Result;
		}

		
	public:
		/**
		 * initializes new SMIMEMessage
		 * @param Input content of SMIME message
		 */
		SMIMEMessage(const std::string &Input) {
			Decrypted = NULL;
			Verified = NULL;

			SSL_library_init();
			SSL_load_error_strings();

			BIO *f = BIO_new(BIO_s_mem());
			BIO_set_mem_eof_return(f, 0);
			if ((unsigned int)BIO_write(f, Input.c_str(), Input.size()) != Input.size()) throw GENERICEXCEPTION("Short write from memory into BIO");
			
			PKCS = SMIME_read_PKCS7(f, &Content);
			if (!PKCS) GENERICSTRMEXCEPTION("Cannot read PKCS#7 information: " << ERR_error_string(ERR_get_error(), NULL));
		}

		/**
		 * decrypts SMIME message with given certificate and private key
		 * @param OwnCertFile path and filename of certificate
		 * @param OwnKeyfile path and filename of key file
		 */
		void Decrypt(const std::string &OwnCertFile, const std::string &OwnKeyFile) {
			FILE *f;

			f = fopen(OwnKeyFile.c_str(), "r");
			if (!f) GENERICSTRMEXCEPTION("Cannot read private key file " << OwnKeyFile);
			EVP_PKEY *Own_Key = PEM_read_PrivateKey(f, NULL, NULL, NULL);
			fclose(f);
			if (!Own_Key) GENERICSTRMEXCEPTION("Cannot parse own private key file.");

			f = fopen(OwnCertFile.c_str(), "r");
			if (!f) GENERICSTRMEXCEPTION("Cannot read own certificate file " << OwnCertFile);
			X509 *Own_Cert = PEM_read_X509(f, NULL, NULL, NULL);
			fclose(f);
			if (!Own_Cert) GENERICSTRMEXCEPTION("Cannot parse own certificate file.");
			
			Decrypted = BIO_new(BIO_s_mem());
			if (PKCS7_decrypt(PKCS, Own_Key, Own_Cert, Decrypted, 0) == 0) {
				GENERICSTRMEXCEPTION("Decryption of PKCS#7 failed: " << ERR_error_string(ERR_get_error(), NULL));
			} else {
				LOGSTRM(Logger::INFO, "Decryption successful.");
			}
		}

		/**
		 * verifies signature of SMIME message
		 * @param CACertificateFile filename and path to CA certificate that should be used for verification
		 */
		bool VerifySignature(const std::string &CACertificateFile) {
			X509_STORE *Store = X509_STORE_new();
			if (!Store) throw GENERICEXCEPTION("Cannot create certificate store.");
			X509_STORE_set_verify_cb_func(Store, NULL);
			
			X509_LOOKUP *Lookup = X509_STORE_add_lookup(Store, X509_LOOKUP_file());
			if (!Lookup) throw GENERICEXCEPTION("Cannot create certificate lookup.");
			if (X509_LOOKUP_load_file(Lookup, CACertificateFile.c_str(), X509_FILETYPE_PEM) != 1) throw GENERICEXCEPTION("Cannot read CA certificate file.");

			Verified = BIO_new(BIO_s_mem());
			int SignatureCheck;
			if ((SignatureCheck = PKCS7_verify(PKCS, NULL, Store, Content, Verified, 0)) != 1) {
				LOGSTRM(Logger::WARNING, "Signature verification failed (" << SignatureCheck << "): " << ERR_error_string(ERR_get_error(), NULL));
				return false;
			} else {
				LOGSTRM(Logger::INFO, "Signature verification passed.");
				return true;
			}
		}

		/**
		 * @returns decrypted message, Decrypt() must have been called!
		 */
		std::string GetDecryptedMessage() {
			if (!Decrypted) throw GENERICEXCEPTION("Cannot retrieve decrypted message if nothing was decrypted.");
			return ReadFromBIO(Decrypted);
		}

		/**
		 * @returns plaintext message if a signature was checked by VerifySignature\
		 */
		std::string GetSignedMessage() {
			if (!Verified) throw GENERICEXCEPTION("Cannot retrieve signed message if no signature was checked.");
			return ReadFromBIO(Verified);
		}
		
		/**
		 * @returns name of certificate that signed message
		 */
		std::string GetSigningParty() {
			if (!Verified) throw GENERICEXCEPTION("Cannot retrieve signing party if no signature was checked.");
			
			STACK_OF(X509) *Signers;
			Signers = PKCS7_get0_signers(PKCS, NULL, 0);
			if (!Signers) throw GENERICEXCEPTION("Cannot extract signing party.");

			X509 *SignCert = sk_X509_value(Signers, 0);
			return X509_NAME_oneline(X509_get_subject_name(SignCert), NULL, 0);
		}
};
#endif

